import argparse
import os
# 导入自定义工具
from utils.h5data import read_h5_data
from utils.plot import *
from scipy import signal

def plot_traceSeq(delta_t, te, traces, ref_trace, t, T_MAX, title=None):
    from pylab import figure, rcParams, np, plt
    from pylab import plot, pcolormesh,scatter
    from pylab import xlim, ylim, xlabel, ylabel,yticks
    from scipy import signal
    rcParams['font.family'] = 'Arial'
    rcParams['font.size'] = '8'

    import matplotlib.gridspec as gridspec
    from utils.math import norm, remove_point_skip
    
    delta_t = remove_point_skip(delta_t, peroid=100)
    te_bjt = te/3600/24+8/24
    YLIM = [te_bjt.min(),te_bjt.max()]
    traces = norm(traces, ONE_AXIS=True)

    plt.close('all')
    fig = figure(figsize=[4,6], dpi=600)
    gs = gridspec.GridSpec(1,4,wspace=0)

    ax1 = fig.add_subplot(gs[0:3])
    X,Y = np.meshgrid(t, te_bjt)
    pcolormesh(X,Y, traces, cmap='bwr', shading='nearest', rasterized=True )  

    plot([T_MAX-1,T_MAX-1],YLIM,lw=0.5, color='gray')
    plot([T_MAX,T_MAX],YLIM,lw=1, color='k')
    plot([T_MAX+1,T_MAX+1],YLIM,lw=0.5, color='gray')

    plot(t, (ref_trace+1)*(YLIM[1]-YLIM[0])/20, lw=0.7, color='k')
    amp = np.abs(signal.hilbert(ref_trace))
    plot(t, (amp+1)*(YLIM[1]-YLIM[0])/20, lw=0.7, color='r')

    # xlim([t.min(),t.max()])
    xlim([T_MAX-3,T_MAX+3])
    ylim([te_bjt.min(),te_bjt.max()])
    xlabel('t(s)')
    ylabel('BJT time')

    ax2 = fig.add_subplot(gs[3:])
    scatter(delta_t,te_bjt, s=5, c='gray', marker='^', edgecolors='none')
    plot(delta_t,te_bjt, '-', color='tab:red', lw=0.5)
    xlim([-50,50])
    ylim([te_bjt.min(),te_bjt.max()])
    xlabel('dt(ms)')
    yticks([])

    plt.suptitle(title)

    return fig, ax1,ax2

def plot_dt_phi(traces, ref_trace, t, 
            te,delta_phi, valid_freq, 
            T_MAX, title=None):
    
    import matplotlib.gridspec as gridspec
    from utils.math import norm, remove_point_skip
    

    te_bjt = te/3600/24+8/24
    ne,nt = traces.shape
    YLIM = [te_bjt.min(),te_bjt.max()]
    # YLIM = [10,15]
    traces = norm(traces, ONE_AXIS=True)

    plt.close('all')
    fig = figure(figsize=[6,6], dpi=600)
    gs = gridspec.GridSpec(1,2,wspace=0)

    ax1 = fig.add_subplot(gs[:,0])
    X,Y = np.meshgrid(t, te_bjt)
    pcolormesh(X,Y, traces, cmap='bwr', shading='nearest', rasterized=True )  

    plot([T_MAX-1,T_MAX-1],YLIM,lw=0.5, color='gray')
    plot([T_MAX,T_MAX],YLIM,lw=1, color='k')
    plot([T_MAX+1,T_MAX+1],YLIM,lw=0.5, color='gray')

    plot(t, (ref_trace+1)*(YLIM[1]-YLIM[0])/20, lw=0.7, color='k')
    amp = np.abs(signal.hilbert(ref_trace))
    plot(t, (amp+1)*(YLIM[1]-YLIM[0])/20, lw=0.7, color='r')

    # xlim([t.min(),t.max()])
    xlim([T_MAX-1,T_MAX+1])
    clim(-traces.max(),traces.max())
    ylim(YLIM)
    xlabel('t(s)')
    ylabel('BJT time')

    ax2 = fig.add_subplot(gs[:,1])
    fs,fe = valid_freq.min(),valid_freq.max()
    # delta_phi_dt = -delta_phi/(2*np.pi*valid_freq)*1000
    delta_phi_dt = delta_phi/np.pi
    X,Y = np.meshgrid(valid_freq, te_bjt)
    pcolormesh(X,Y, delta_phi_dt, cmap='bwr', shading='nearest', rasterized=True )  
    
    spec = np.fft.fft(traces,axis=-1).sum(axis=0)[:int(nt//2)]
    spec = np.abs(spec)
    freqs = np.fft.fftfreq(traces.shape[1], d=t[1]-t[0])[:int(nt//2)]
    spec = spec/spec.max()*4
    plot(freqs, spec+YLIM[0], lw=1,color='k')

    xlim([fs,fe])
    # xlim([5,15])
    ylim(YLIM)
    xlabel('f(Hz)')
    yticks([])
    # clim(-delta_phi_dt.std()*2,delta_phi_dt.std()*2)
    clim(-0.2,0.2)
    colorbar()

    plt.suptitle(title)

    return fig, ax1,ax2


if __name__ == '__main__':

    

    parser = argparse.ArgumentParser(description='Plot traceSeq h5 data')
    parser.add_argument('-input', default='', help='input h5 file')
    parser.add_argument('-figroot', default='figures/7.dt.change.figures', help='root to save figs')
    # parser.add_argument('-figroot', default='figures/debug', help='root to save figs')
    args = parser.parse_args()
    INPUT_FILE = args.input
    FIG_ROOT = args.figroot

    # 读取元数据
    datasets, args_in_file = read_h5_data(INPUT_FILE, 
                                             ['MARKER','EMARKER','DATE','BASE_STAT'],
                                             group_name='metadata',
                                             read_attrs=True)
    MARKER = datasets[0].decode() 
    EMARKER = datasets[1].decode() 
    DATE = datasets[2].decode()
    SOURCE = datasets[3].decode()

    # 读取数据
    R, delta_t, te, traces, ref, t, T_MAX, r,d_phi,freqs = read_h5_data(INPUT_FILE,
                    ['R','delta_t','te','traces','ref','t','T_MAX','r',
                     'd_phi','freqs'])
    
    # 解码台站名
    R = [r.decode() if isinstance(r, bytes) else r for r in R]

    FIG_ROOT = f'{FIG_ROOT}/{DATE}.{EMARKER}.{MARKER}'
    if not os.path.exists(FIG_ROOT):
        os.makedirs(FIG_ROOT)

    # 为每个接收器绘图
    for i, name in enumerate(R):
        # 提取第i个接收器的数据
        delta_t_i = delta_t[i, :]
        traces_i = traces[i, :, :]
        ref_i = ref[i, :]
        T_MAX_i = T_MAX[i]
        r_i = r[i]
        d_phi_i = d_phi[i, :]
        
        # fig,_,_ = plot_traceSeq(delta_t_i, te, traces_i, ref_i, t, T_MAX_i, 
        #                        title=f'{SOURCE}-{name}.T={T_MAX_i:.1f}.V={r_i/T_MAX_i:.2f}')

        fig,_,_ = plot_dt_phi(traces_i, ref_i, t, 
                              te,d_phi_i, freqs, 
                              T_MAX_i, 
                               title=f'{SOURCE}-{name}.T={T_MAX_i:.1f}.V={r_i/T_MAX_i:.2f}')

        output_file = f'{FIG_ROOT}/{name}.{MARKER}.png'
        print(output_file)
        fig.tight_layout()
        fig.savefig(output_file)